/*
* This file is a part of the open-eBackup project.
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
* If a copy of the MPL was not distributed with this file, You can obtain one at
* http://mozilla.org/MPL/2.0/.
*
* Copyright (c) [2024] Huawei Technologies Co.,Ltd.
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
*/
#ifdef WIN32
#include "HyperVVolumeHandler.h"
#include "protect_engines/hyperv/api/powershell/model/GetVMVolumesModel.h"
#include "Intsafe.h"


const int CHUNK_SIZE = 16384; // 16384 bytes
const UINT64 GB = 1024 * 1024 * 1024;
const UINT64 SEGMENT_SIZE_PER_GB = GB / (CHUNK_SIZE * 2);

std::shared_ptr<uint8_t[]> VirtPlugin::VolumeHandler::m_allZeroDirtyRangeDataPtr = nullptr;

namespace HyperVPlugin {

int32_t HyperVVolumeHandler::InitializeVolumeInfo()
{
    m_ranges = std::make_unique<QUERY_CHANGES_VIRTUAL_DISK_RANGE[]>(SEGMENT_SIZE_PER_GB);
    if (m_ranges == nullptr) {
        ERRLOG("m_ranges is null.");
        return FAILED;
    }
    if (m_jobHandle == nullptr) {
        ERRLOG("Job handler is null.");
        return FAILED;
    }
    return SUCCESS;
}

bool HyperVVolumeHandler::OpenDisk(const std::string &diskPath, HANDLE &volHandler)
{
    VIRTUAL_STORAGE_TYPE storageType;
    storageType.DeviceId = VIRTUAL_STORAGE_TYPE_DEVICE_UNKNOWN;
    OPEN_VIRTUAL_DISK_PARAMETERS params;
    memset_s(&params, sizeof(params), 0, sizeof(params));
    params.Version = OPEN_VIRTUAL_DISK_VERSION_2;
    params.Version2.GetInfoOnly = FALSE;
    DWORD ret = OpenVirtualDisk(&storageType, String2WstringByUtf8(diskPath).c_str(),
        VIRTUAL_DISK_ACCESS_NONE, OPEN_VIRTUAL_DISK_FLAG_NONE, &params, &volHandler);
    if (ret != ERROR_SUCCESS || volHandler == INVALID_HANDLE_VALUE) {
        ERRLOG("OpenVirtualDisk  failed, ret = %d, error code: %d.", ret, GetLastError());
        return false;
    }
    return true;
}

bool HyperVVolumeHandler::AttachDisk(HANDLE &volHandler, ATTACH_VIRTUAL_DISK_FLAG attachFlags)
{
    ATTACH_VIRTUAL_DISK_PARAMETERS attachParameters;
    memset_s(&attachParameters, sizeof(attachParameters), 0, sizeof(attachParameters));
    attachParameters.Version = ATTACH_VIRTUAL_DISK_VERSION_1;
    DWORD ret = AttachVirtualDisk(volHandler, NULL, attachFlags, 0, &attachParameters, NULL);
    if (ret != ERROR_SUCCESS && ret != ERROR_SHARING_VIOLATION) {
        ERRLOG("AttachVirtualDisk failed, ret = %d, error code: %d.", ret, GetLastError());
        return false;
    }
    return true;
}

int32_t HyperVVolumeHandler::Open(const VolOpenMode &mode)
{
    // VolOpenMode::READ_WRITE path is for restore
    if (mode != VolOpenMode::READ_WRITE) {
        // find absolute path of target disk file
        ERRLOG("Unexpected mode: %d", mode);
        return FAILED;
    }
    m_IfRestore = true;
    // Find target disk file path from target volumne info
    Json::Value targetDisk;
    if (!Module::JsonHelper::JsonStringToJsonValue(m_volInfo.m_extendInfo, targetDisk)) {
        ERRLOG("Convert %s failed", WIPE_SENSITIVE(m_volInfo.m_extendInfo).c_str());
        return FAILED;
    }
    VolumeInfo targetDiskObj;
    if (!Module::JsonHelper::JsonValueToStruct(targetDisk, targetDiskObj)) {
        ERRLOG(("Turning jsonValue to struct failed, json: %s", targetDisk.toStyledString()).c_str());
        return FAILED;
    }
    INFOLOG("Target volume file path: %s", targetDiskObj.m_path.c_str());
    if (!OpenDisk(targetDiskObj.m_path, m_writeVhdHandle)) {
        ERRLOG("Opendisk %s failed.", targetDiskObj.m_path.c_str());
        return FAILED;
    }
    ATTACH_VIRTUAL_DISK_FLAG attachFlags = ATTACH_VIRTUAL_DISK_FLAG_NO_LOCAL_HOST;
    if (!AttachDisk(m_writeVhdHandle, attachFlags)) {
        ERRLOG("Hyperv attach disk(%s) failed.", targetDiskObj.m_path.c_str());
        return FAILED;
    }
    return SUCCESS;
}

/**
 * @brief 写卷
 *
 * @param offsetInBytes -- where to write in target file
 * @param bufSizeInBytes -- how long to write in target file
 * @param buf -- where the data is
 * @return int32_t
 */

int32_t HyperVVolumeHandler::WriteBlocks(const uint64_t &offsetInBytes, uint64_t &bufSizeInBytes,
                                         std::shared_ptr<uint8_t[]> &buf)
{
    LARGE_INTEGER offset;
    offset.QuadPart = offsetInBytes;
    OVERLAPPED overlapped;
    memset_s(&overlapped, sizeof(overlapped), 0, sizeof(overlapped));
    overlapped.Offset = offset.LowPart;
    overlapped.OffsetHigh = offset.HighPart;
    DWORD bytesWritten;
    INFOLOG("BEGIN READBLOCKS offset %d, offsetHight %d.", overlapped.Offset, overlapped.OffsetHigh);
    BOOL ret = WriteFile(m_writeVhdHandle, (LPVOID)buf.get(), bufSizeInBytes, &bytesWritten, &overlapped);
    if (ret == TRUE) {
        DBGLOG("WriteFile success, %d.", bytesWritten);
        return bytesWritten;
    }
    ret = GetLastError();
    if (ret != ERROR_IO_PENDING) {
        ERRLOG("Read file failed, error: %d.", ret);
        return 0;
    }
    if (!GetOverlappedResult(m_writeVhdHandle, &overlapped, &bytesWritten, TRUE)) {
        ERRLOG("Get over lapped failed, %llu, %llu, ret: %d.", offsetInBytes, bufSizeInBytes, GetLastError());
        return FAILED;
    }
    return SUCCESS;
}

/**
 * @brief 打开卷
 *
 * @param mode
 * @param jobInfo 任务信息，包括前一次快照和当前快照信息
 * @return int32_t
 */
int32_t HyperVVolumeHandler::Open(const VolOpenMode &mode, const BackupSubJobInfo &jobInfo)
{
    std::string preDiskPath = jobInfo.m_preSnapshotInfo.m_snapshotWwn;
    std::string curDiskPath = jobInfo.m_curSnapshotInfo.m_snapshotWwn;
    std::string diskPath;
    if (preDiskPath.empty()) {
        WARNLOG("preDiskPath.empty(), maybe new disk for Incre backup. Check curDiskPath.");
        if (curDiskPath.empty()) {
            ERRLOG("curDiskPath.empty(). cannot get disk path");
            return FAILED;
        }
        diskPath = curDiskPath;
    } else {
        diskPath = preDiskPath;
    }
    INFOLOG("Disk path is %s.", diskPath.c_str());
    if (!OpenDisk(diskPath, m_readVhdHandle)) {
        ERRLOG("Hyperv open disk(%s) failed.", diskPath.c_str());
        return FAILED;
    }
    ATTACH_VIRTUAL_DISK_FLAG attachFlags = ATTACH_VIRTUAL_DISK_FLAG_NO_LOCAL_HOST | ATTACH_VIRTUAL_DISK_FLAG_READ_ONLY;
    if (!AttachDisk(m_readVhdHandle, attachFlags)) {
        ERRLOG("Hyperv attach disk(%s) failed.", diskPath.c_str());
        return FAILED;
    }

    INFOLOG("Open %s success.", diskPath.c_str());
    return SUCCESS;
}


int32_t HyperVVolumeHandler::Close()
{
    if (m_IfRestore) {
        INFOLOG("Close for restore job");
        if (m_writeVhdHandle != INVALID_HANDLE_VALUE) {
            CloseHandle(m_writeVhdHandle);
        }
        return SUCCESS;
    }
    if (m_readVhdHandle != INVALID_HANDLE_VALUE) {
        CloseHandle(m_readVhdHandle);
    }
    return SUCCESS;
}

int32_t HyperVVolumeHandler::ReadBlocks(const uint64_t& offsetInBytes, uint64_t& bufferSizeInBytes,
    std::shared_ptr<uint8_t[]> &buffer, std::shared_ptr<uint8_t[]> &calBuffer, std::shared_ptr<uint8_t[]> &readBuffer)
{
    if (offsetInBytes + bufferSizeInBytes > m_volInfo.m_volSizeInBytes) {
        bufferSizeInBytes = m_volInfo.m_volSizeInBytes - offsetInBytes;
        INFOLOG("Read len (%llu, %llu, %llu).", offsetInBytes, bufferSizeInBytes, m_volInfo.m_volSizeInBytes);
    }
    std::lock_guard<std::mutex> lock(m_repoMutex);
    LARGE_INTEGER offset;
    offset.QuadPart = offsetInBytes;
    OVERLAPPED overlapped;
    memset_s(&overlapped, sizeof(overlapped), 0, sizeof(overlapped));
    overlapped.Offset = offset.LowPart;
    overlapped.OffsetHigh = offset.HighPart;
    DWORD bytesRead;
    DBGLOG("BEGIN READBLOCKS offset %d, offsetHight %d.", overlapped.Offset, overlapped.OffsetHigh);
    BOOL ret = ReadFile(m_readVhdHandle, (LPVOID)buffer.get(), bufferSizeInBytes, &bytesRead, &overlapped);
    if (ret == TRUE) {
        if (memcmp(buffer.get(), GetAllZeroDirtyRangeDataPtr().get(), bufferSizeInBytes) == 0) {
            return DATA_ALL_ZERO_IGNORE_WRITE;
        }
        return SUCCESS;
    }
    ret = GetLastError();
    if (ret != ERROR_IO_PENDING) {
        ERRLOG("Read file failed, error: %d.", ret);
        return FAILED;
    }
    if (!GetOverlappedResult(m_readVhdHandle, &overlapped, &bytesRead, TRUE)) {
        ERRLOG("Get over lapped failed, %llu, %llu, ret: %d.", offsetInBytes, bufferSizeInBytes, GetLastError());
        return FAILED;
    }
    if (memcmp(buffer.get(), GetAllZeroDirtyRangeDataPtr().get(), bufferSizeInBytes) == 0) {
        DBGLOG("Get zero ignore write");
        return DATA_ALL_ZERO_IGNORE_WRITE;
    }
    return SUCCESS;
}

bool HyperVVolumeHandler::AddDirtyRanges(DirtyRanges &dirtyRanges, ULONG rangeCount)
{
    for (ULONG i = 0; i < rangeCount; ++i) {
        DBGLOG("Offset: %lld, length: %lld, reserved: %lld", m_ranges[i].ByteOffset, m_ranges[i].ByteLength,
            m_ranges[i].Reserved);
        if (!dirtyRanges.AddRange(DirtyRange(m_ranges[i].ByteOffset, m_ranges[i].ByteLength))) {
            ERRLOG("Add m_ranges failed.");
            return false;
        }
    }
    return true;
}

/**
* @brief 将磁盘按GE分小段计算差异位图
*
* @param rctIdStr
* @param diskSize 磁盘大小, 单位字节
* @param dirtyRanges
* @return int32_t
*/
int32_t HyperVVolumeHandler::GetDirtyRangesFromVirtDisk(const std::string &rctIdStr, uint64_t startOffset,
    uint64_t endOffset, DirtyRanges &dirtyRanges)
{
    if (rctIdStr.empty()) {
        ERRLOG("Rct id is empty.");
        return FAILED;
    }
    PCWSTR rctId = _com_util::ConvertStringToBSTR(rctIdStr.c_str());
    ULONG64 processedLen = 0;
    uint64_t byteLength = 0;
    uint64_t totRangeCount = 0;
    for (int i = startOffset / GB; i < endOffset / GB; i++) {
        byteLength = (i + 1) * GB > endOffset ? endOffset - i * GB : GB;
        ULONG rangeCount = SEGMENT_SIZE_PER_GB;
        DWORD ret = QueryChangesVirtualDisk(m_readVhdHandle, rctId, i * GB, byteLength,
            QUERY_CHANGES_VIRTUAL_DISK_FLAG_NONE, m_ranges.get(), &rangeCount, &processedLen);
        if (ret != ERROR_SUCCESS) {
            ERRLOG("QueryChangesVirtualDisk(%llu, %llu) failed.", startOffset, endOffset);
            return FAILED;
        }
        DBGLOG("Begin to query changes %d, %llu, %llu, range count(%u), processedLen(%llu).", i, i * GB,
            byteLength, rangeCount, processedLen);
        if (!AddDirtyRanges(dirtyRanges, rangeCount)) {
            ERRLOG("Add ranges failed.");
            return FAILED;
        }
        totRangeCount += rangeCount;
    }
    return SUCCESS;
}

int32_t HyperVVolumeHandler::GetDirtyRanges(const VolSnapInfo &preVolSnapshot, const VolSnapInfo &curVolSnapshot,
    DirtyRanges &dirtyRanges, const uint64_t startOffset, uint64_t &endOffset)
{
    INFOLOG("PreVolSnapshot.m_extendInfo %s.", preVolSnapshot.m_extendInfo.c_str());
    if (!preVolSnapshot.m_extendInfo.empty()) {
        Json::Value extendInfo;
        if (!Module::JsonHelper::JsonStringToJsonValue(preVolSnapshot.m_extendInfo, extendInfo)) {
            return FAILED;
        }
        std::string rctId = extendInfo["RctId"].asString();
        INFOLOG("Rct id is %s.", rctId.c_str());
        int ret = GetDirtyRangesFromVirtDisk(rctId, startOffset, endOffset, dirtyRanges);
        if (ret == SUCCESS) {
            return SUCCESS;
        }
        WARNLOG("GetDirtyRangesFromVirtDisk failed.");
    }
    INFOLOG("Volume(%s) Use full ranges(%lld, %lld).", curVolSnapshot.m_volUuid.c_str(), startOffset, endOffset);
    dirtyRanges.UseFull();
    DirtyRange everything(startOffset, endOffset - startOffset);
    (void)dirtyRanges.AddRange(everything);
    return SUCCESS;
}

uint64_t HyperVVolumeHandler::GetVolumeSize()
{
    ERRLOG("not realized, do not use this mothod");
    return 0;
}

int32_t HyperVVolumeHandler::TestDeviceConnection(const std::string &authExtendInfo, int32_t &erroCode)
{
    ERRLOG("not realized, do not use this mothod");
    return SUCCESS;
}

int32_t HyperVVolumeHandler::CleanLeftovers()
{
    ERRLOG("not realized, do not use this mothod");
    return SUCCESS;
}

int32_t HyperVVolumeHandler::Flush()
{
    ERRLOG("not realized, do not use this mothod");
    return SUCCESS;
}
}
#endif